'''
历史消息裁剪策略示例

问题背景：
历史消息越多，对大模型token消耗越大，同时可能带来：
1. 上下文长度超限错误
2. 响应速度变慢
3. 成本增加
4. 无关历史干扰当前对话

解决方案：消息裁剪（Trimming）
保留最近N条消息，删除早期消息，确保上下文在可控范围内
'''
import os

from dotenv import load_dotenv
from langchain_community.chat_message_histories import ChatMessageHistory
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables import RunnableWithMessageHistory, RunnablePassthrough
from langchain_openai import ChatOpenAI

# 自动加载 .env 文件中的环境变量
load_dotenv()

# ========== 模拟构建历史消息记录 ==========
'''
创建模拟对话历史，展示裁剪前后的变化
'''
temp_chat_history = ChatMessageHistory()

# 添加6轮对话，模拟真实的聊天场景
temp_chat_history.add_user_message("hello,我叫未来")
temp_chat_history.add_ai_message("你好")
temp_chat_history.add_user_message("今天天气怎么样")
temp_chat_history.add_ai_message("今天天气晴朗，适合出去玩")
temp_chat_history.add_user_message("那我准备出门")
temp_chat_history.add_ai_message("好的，祝你玩的愉快")

print("=== 初始历史消息 ===")
for i, message in enumerate(temp_chat_history.messages):
    print(f"{i+1}. {message.type}: {message.content}")

# ========== 构建基础对话系统 ==========
prompt = ChatPromptTemplate.from_messages([
    ("system", "你是一个乐于助人的助手。尽力回答所有问题，提供的聊天历史包括与您交谈的用户的记录。"),
    MessagesPlaceholder(variable_name="chat_history"),  # 历史消息注入点
    ("human", "{input}")  # 当前用户输入
])

# 模型选择说明：
# - Qwen3-VL: 支持视觉和语言的多模态模型，适合复杂理解任务
# - Qwen3-Coder: 专为代码生成优化的模型，适合SQL生成场景
model_name = ""  # 235B总参数，22B激活参数
# model_name = "Qwen/Qwen3-Coder-480B-A35B-Instruct"  # 代码专用模型


# 初始化大语言模型
llm = ChatOpenAI(api_key=os.getenv("MODELSCOPE_API_KEY"), openai_api_base=os.getenv("MODELSCOPE_API_BASE"), model="Qwen/Qwen3-32B", streaming=True, temperature=0)

# 构建基础处理链：提示模板 -> 大模型
chain = prompt | llm

# ========== 消息裁剪功能 ==========
def trim_message(chain_input):
    """
    消息裁剪函数：保留最近N条消息，删除早期历史

    裁剪策略说明：
    - 保留最近2条消息（最新对话）
    - 删除所有早期消息
    - 适用于关注最近对话的场景

    应用场景：
    1. 实时对话系统 - 只关心最近互动
    2. 上下文严格受限 - 必须控制token数量
    3. 话题快速切换 - 避免历史干扰

    优缺点：
    ✅ 优点：简单高效，严格控token，避免历史干扰
    ❌ 缺点：丢失长期记忆，可能遗忘重要信息
    """
    stored_messages = temp_chat_history.messages

    # 检查消息数量，小于等于2条不需要裁剪
    if len(stored_messages) <= 2:
        print("📊 当前消息数 ≤ 2，无需裁剪")
        return False

    print(f"✂️ 开始裁剪消息...")
    print(f"   原始消息数: {len(stored_messages)}")
    print(f"   保留最近: 2 条")

    # 显示将被删除的早期消息
    print("   🗑️ 删除的早期消息:")
    for i, message in enumerate(stored_messages[:-2]):
        print(f"     - {message.type}: {message.content}")

    # 显示将保留的最近消息
    print("   💾 保留的最近消息:")
    for i, message in enumerate(stored_messages[-2:]):
        print(f"     - {message.type}: {message.content}")

    # 执行裁剪：清空所有消息，只保留最近2条
    temp_chat_history.clear()
    for message in stored_messages[-2:]:
        temp_chat_history.add_message(message)

    print(f"✅ 消息裁剪完成，当前剩余 {len(temp_chat_history.messages)} 条消息")
    return True

# ========== 构建带历史管理的对话链 ==========
'''
RunnableWithMessageHistory 自动管理对话历史：
- 根据session_id获取对应历史
- 自动将历史注入到提示词中
- 维护对话的连续性
'''
chain_with_message_history = RunnableWithMessageHistory(
    chain,  # 基础对话链
    lambda session_id: temp_chat_history,  # 历史消息提供函数
    input_messages_key="input",           # 用户输入的key
    history_messages_key="chat_history",  # 历史消息在提示词中的key
)

# ========== 构建带裁剪功能的完整链 ==========
'''
chain_with_trimming 的工作流程：
1. RunnablePassthrough.assign(messages_trimmed=trim_message)
   - 先执行 trim_message 函数进行消息裁剪
   - 将返回值赋给 messages_trimmed 变量
   - 保持原始输入继续传递

2. | chain_with_message_history
   - 将处理后的输入传递给带历史管理的对话链
   - 使用裁剪后的历史生成回复

设计理念：先裁剪，再对话，确保每次对话都在可控的上下文范围内
'''
chain_with_trimming = (
    # 先执行消息裁剪
    RunnablePassthrough.assign(messages_trimmed=trim_message)
    # 再执行带历史管理的对话
    | chain_with_message_history
)

# ========== 测试演示 ==========
print("\n" + "="*60)
print("🚀 开始测试消息裁剪功能")
print("="*60)

# 测试问题1：依赖早期历史的问题（裁剪后可能无法回答）
test_question_1 = "我今天出门了吗？"
print(f"\n📝 测试问题1: {test_question_1}")
print("💡 这个问题依赖早期历史：用户说『那我准备出门』，AI回复『好的，祝你玩的愉快』")

response_1 = chain_with_trimming.invoke(
    {"input": test_question_1},
    config={"configurable": {"session_id": "lsq"}}
)
print(f"🤖 AI回复: {response_1.content}")

# 显示裁剪后的历史状态
print("\n📋 裁剪后的历史记录:")
for i, message in enumerate(temp_chat_history.messages):
    print(f"   {i+1}. {message.type}: {message.content}")

# ========== 进阶演示：不同裁剪策略对比 ==========
def demonstrate_trimming_strategies():
    """
    演示不同的消息裁剪策略
    """
    print("\n" + "="*60)
    print("🔬 进阶演示：不同裁剪策略对比")
    print("="*60)

    # 策略1：固定数量裁剪（当前实现）
    print("\n🎯 策略1: 固定数量裁剪")
    print("   - 保留最近N条消息")
    print("   - 简单直接，严格控token")
    print("   - 可能丢失重要早期信息")

    # 策略2：智能裁剪（按重要性）
    print("\n🎯 策略2: 智能重要性裁剪")
    print("   - 保留包含关键信息（姓名、偏好、决策）的消息")
    print("   - 删除日常寒暄和无关对话")
    print("   - 需要额外的重要性判断逻辑")

    # 策略3：时间窗口裁剪
    print("\n🎯 策略3: 时间窗口裁剪")
    print("   - 保留最近X分钟内的消息")
    print("   - 适合实时性强的场景")
    print("   - 需要记录消息时间戳")

    # 策略4：Token数量裁剪
    print("\n🎯 策略4: Token数量裁剪")
    print("   - 计算累计token数，超限时裁剪")
    print("   - 最精确的上下文控制")
    print("   - 需要token计数功能")

def create_adaptive_trimmer(max_messages=3, keep_important=True):
    """
    创建自适应裁剪器示例
    """
    def adaptive_trim(chain_input):
        stored_messages = temp_chat_history.messages

        if len(stored_messages) <= max_messages:
            return False

        # 简单的重要性判断：包含关键信息的消息
        important_keywords = ['我叫', '名字是', '计划', '项目', '重要']
        important_messages = []
        other_messages = []

        for msg in stored_messages:
            content = msg.content
            is_important = any(keyword in content for keyword in important_keywords)
            if is_important and keep_important:
                important_messages.append(msg)
            else:
                other_messages.append(msg)

        # 保留重要消息 + 最近的普通消息
        messages_to_keep = important_messages + other_messages[-max_messages:]

        temp_chat_history.clear()
        for msg in messages_to_keep[-max_messages*2:]:  # 限制总数
            temp_chat_history.add_message(msg)

        print(f"🔧 自适应裁剪完成，保留 {len(temp_chat_history.messages)} 条消息")
        return True

    return adaptive_trim

# 取消注释运行进阶演示
# demonstrate_trimming_strategies()

# 测试自适应裁剪器
# adaptive_trimmer = create_adaptive_trimmer(max_messages=2)
# chain_with_adaptive_trimming = (
#     RunnablePassthrough.assign(messages_trimmed=adaptive_trimmer)
#     | chain_with_message_history
# )

# ========== 裁剪 vs 汇总 对比 ==========
'''
📊 消息裁剪 vs 消息汇总 对比：

消息裁剪（当前示例）：
✅ 优点：
   - 实现简单，计算开销小
   - 严格保证上下文长度
   - 保留原始对话内容
❌ 缺点：
   - 直接丢失被裁剪的信息
   - 可能遗忘重要长期记忆
   - 对话连续性可能中断

消息汇总（前一个示例）：
✅ 优点：
   - 保留所有信息的精华
   - 维持长期记忆能力
   - 对话连续性更好
❌ 缺点：
   - 实现复杂，需要额外LLM调用
   - 摘要可能失真或遗漏细节
   - 增加成本和延迟

选择建议：
- 实时对话、严格token限制 → 消息裁剪
- 长期记忆、重要信息保留 → 消息汇总
- 混合策略 → 重要信息保留 + 普通消息裁剪
'''

print("\n" + "="*60)
print("🎯 总结：消息裁剪的核心价值")
print("="*60)
print("1. 🚀 性能优化：严格控制token消耗，提升响应速度")
print("2. 💰 成本控制：减少API调用费用")
print("3. 🎯 焦点保持：避免无关历史干扰当前对话")
print("4. 🔧 简单可靠：实现和维护成本低")
print("5. ⚡ 实时适用：适合需要快速响应的对话场景")